import math
import numpy
import window
import stft
import filter_bank
from scipy.io import wavfile


class audioFeatureExtractor:
    def __init__(self, **kwargs):
        Window = None
        OverlapLength = 512
        SampleRate = 44100
        SpectralDescriptorInput = 0
        FFTLength = -1
        FilterBank = None
        for key in kwargs:
            v = kwargs[key]
            if key == 'SampleRate':
                SampleRate = int(v)
            elif key == 'FFTLength':
                FFTLength = int(v)
            elif key == 'Window':
                Window = v
            elif key == 'OverlapLength':
                OverlapLength = int(v)
            elif key == 'FilterBank':
                FilterBank = v
        if Window is None:
            Window = window.hamming(1024, 'periodic')
        if FFTLength < 0:
            FFTLength = len(Window)

        self.SampleRate = SampleRate
        self.FFTLength = FFTLength
        self.Window = Window
        self.OverlapLength = OverlapLength
        self.FilterBank = FilterBank

    def extract(self, x):
        # stft
        s, _, _ = stft.stft(x, self.SampleRate, Window=self.Window, overlapLength=self.OverlapLength,
                            FFTLength=self.FFTLength, FrequencyRange='onesided')

        # magnitude
        mag_frames = numpy.absolute(s)
        # Power Spectrum
        pow_frames = ((1.0 / self.FFTLength) * ((mag_frames) ** 2))

        # apply filter bank
        result = filter_bank.apply_filter_bank(pow_frames, self.FilterBank)
        return result



if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import sys
    from scipy.io import wavfile

    if len(sys.argv)>1:
        fn = sys.argv[1]
    else:
        fn = "stft.wav"

    # Known sample rate of the data set.
    fs = 16e3;

    fs, data = wavfile.read(fn)

    x = data / 32768

    segmentDuration = 1;
    frameDuration = 0.025;
    hopDuration = 0.010;

    FFTLength = 512;
    numBands = 50;

    segmentSamples = round(segmentDuration * fs)
    frameSamples = round(frameDuration * fs)
    hopSamples = round(hopDuration * fs)
    overlapSamples = frameSamples - hopSamples

    bark_fbank, _, _ = filter_bank.gen_bark_filter_bank(FFTLength, fs, numBands)
    afe = audioFeatureExtractor(
        SampleRate=fs,
        FFTLength=FFTLength,
        Window=window.hann(frameSamples, "periodic"),
        OverlapLength=overlapSamples,
        FilterBank=bark_fbank)

    bark_spectrum = afe.extract(x)
    x=bark_spectrum.shape[0]
    y=bark_spectrum.shape[1]
    Y = numpy.arange(0, y/100, 0.01)
    X = numpy.arange(0, x/100, 0.01)
    c = plt.pcolormesh(X, Y, bark_spectrum.T, shading='nearest')

    plt.colorbar(c)
    plt.ylabel('Frequency [KHz]')
    plt.xlabel('Time [sec]')
    plt.show()
